#!/usr/bin/env python3
"""
RAG + 本地模型
"""

import os
import logging
from typing import List, Dict, Any
from dotenv import load_dotenv

# LangChain 相关导入
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import WebBaseLoader, TextLoader
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate
from langchain.schema import Document

from simple.ollama_llm import OllamaLLM

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

prompt_template = """
基于以下文档内容，回答用户问题。
文档内容：
{context}

用户问题：{question}

回答指导：
- 只基于提供的文档内容回答
- 如果文档中没有相关信息，明确说明
- 回答要准确、简洁、有条理
- 使用中文回答

要求：直接给出最终答案，不要说"首先...然后...最后"或"让我想想"等过程性表述。不输出<think></think>之间的内容

回答：
"""

class RAGSystem:
    def __init__(self, model_name: str):
        """初始化系统"""
        load_dotenv()

        # 确保 USER_AGENT 环境变量已设置
        if not os.environ.get('USER_AGENT'):
            os.environ['USER_AGENT'] = 'Mozilla/5.0 (compatible; RAG-System/1.0)'
            logger.info("已自动设置 USER_AGENT 环境变量")

        self.model_name = model_name
        self.embeddings = None
        self.vectorstore = None
        self.llm = None
        self.qa_chain = None

        logger.info(f"RAG 系统初始化完成，将使用本地模型: {model_name}")

    def setup_embeddings(self):
        """设置嵌入模型"""
        logger.info("正在加载嵌入模型...")

        try:
            # 使用轻量级的多语言嵌入模型
            self.embeddings = HuggingFaceEmbeddings(
                model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
                model_kwargs={'device': 'mps'},
                encode_kwargs={'normalize_embeddings': True}
            )
            logger.info("嵌入模型加载完成")
        except Exception as e:
            logger.error(f"嵌入模型加载失败: {str(e)}")
            logger.info("如果是网络问题，可以尝试使用其他嵌入模型或离线模型")
            raise

    def setup_llm(self):
        """设置本地语言模型"""
        logger.info(f"正在配置本地模型: {self.model_name}")

        # 使用完全修复版本的 Ollama LLM
        self.llm = OllamaLLM(
            model_name=self.model_name,
            base_url="http://localhost:11434",
            temperature=0.7,
            max_tokens=1000
        )

        # 测试连接
        try:
            # 先验证 api_url 属性是否可以正常访问
            logger.info(f"测试 api_url 属性访问: {self.llm.api_url}")

            test_response = self.llm._call("你好")
            if test_response and not any(error in test_response for error in
                                         ["API 调用失败", "网络请求失败", "处理响应时出错", "属性访问错误"]):
                logger.info(f"本地模型连接成功，测试响应: {test_response[:50]}...")
            else:
                raise Exception(f"模型响应异常: {test_response}")
        except Exception as e:
            logger.error(f"本地模型连接失败: {str(e)}")
            self._print_ollama_troubleshooting()
            raise

    def _print_ollama_troubleshooting(self):
        """打印 Ollama 故障排除信息"""
        logger.info("=== Ollama 故障排除 ===")
        logger.info("1. 确保 Ollama 服务正在运行: ollama serve")
        logger.info(f"2. 确保模型已下载: ollama pull {self.model_name}")
        logger.info("3. 检查服务状态: curl http://localhost:11434/api/tags")
        logger.info("4. 检查端口是否被占用: netstat -an | grep 11434")
        logger.info("5. 如果仍有问题，尝试重启 Ollama 服务")

    def load_documents(self, sources: List[str]) -> List[Document]:
        """加载文档（完全修复版本）"""
        documents = []

        # 确保 sources 是列表类型
        if isinstance(sources, str):
            logger.warning(f"收到字符串类型的源，转换为列表: {sources}")
            sources = [sources]

        logger.info(f"准备加载 {len(sources)} 个文档源")

        for i, source in enumerate(sources):
            logger.info(f"正在处理第 {i + 1}/{len(sources)} 个源: {source}")

            try:
                if source.startswith(('http://', 'https://')):
                    # 加载网页内容
                    loader = WebBaseLoader(
                        source,
                        header_template={
                            'User-Agent': os.environ.get('USER_AGENT', 'Mozilla/5.0 (compatible; RAG-System/1.0)')
                        }
                    )
                    docs = loader.load()
                    logger.info(f"成功加载网页: {source}，包含 {len(docs)} 个文档")
                else:
                    # 加载本地文件
                    if not os.path.exists(source):
                        logger.warning(f"文件不存在，跳过: {source}")
                        continue

                    loader = TextLoader(source, encoding='utf-8')
                    docs = loader.load()
                    logger.info(f"成功加载文件: {source}，包含 {len(docs)} 个文档")

                documents.extend(docs)

            except Exception as e:
                logger.error(f"加载文档失败 {source}: {str(e)}")

        logger.info(f"总共成功加载 {len(documents)} 个文档")
        return documents

    def split_documents(self, documents: List[Document]) -> List[Document]:
        """分割文档为小块"""
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=800,
            chunk_overlap=150,
            length_function=len,
            separators=["\n\n", "\n", "。", "！", "？", "；", " ", ""]
        )

        split_docs = text_splitter.split_documents(documents)
        logger.info(f"文档分割完成，共生成 {len(split_docs)} 个文档块")
        return split_docs

    def create_vectorstore(self, documents: List[Document]):
        """创建向量数据库"""
        if not self.embeddings:
            raise ValueError("请先初始化嵌入模型")

        logger.info("正在创建向量数据库...")

        try:
            self.vectorstore = Chroma.from_documents(
                documents=documents,
                embedding=self.embeddings,
                persist_directory="../tmp",
                collection_name="rag_tmp"
            )

            self.vectorstore.persist()
            logger.info("向量数据库创建完成")
        except Exception as e:
            logger.error(f"向量数据库创建失败: {str(e)}")
            raise

    def setup_qa_chain(self):
        """设置问答链"""
        if not self.vectorstore or not self.llm:
            raise ValueError("请先设置向量数据库和语言模型")

        PROMPT = PromptTemplate(
            template=prompt_template,
            input_variables=["context", "question"]
        )

        self.qa_chain = RetrievalQA.from_chain_type(
            llm=self.llm,
            chain_type="stuff",
            retriever=self.vectorstore.as_retriever(
                search_type="similarity",
                search_kwargs={"k": 3}
            ),
            chain_type_kwargs={"prompt": PROMPT},
            return_source_documents=True
        )

        logger.info("问答链设置完成")

    def query(self, question: str) -> Dict[str, Any]:
        """查询问答系统"""
        if not self.qa_chain:
            raise ValueError("请先设置问答链")

        logger.info(f"正在处理问题: {question}")

        try:
            result = self.qa_chain.invoke({"query": question})

            response = {
                "question": question,
                "answer": result["result"],
                "source_documents": [
                    {
                        "content": doc.page_content[:150] + "...",
                        "metadata": doc.metadata
                    }
                    for doc in result["source_documents"]
                ]
            }

            logger.info("问题处理完成")
            return response

        except Exception as e:
            logger.error(f"查询失败: {str(e)}")
            return {
                "question": question,
                "answer": f"抱歉，处理问题时出现错误: {str(e)}",
                "source_documents": []
            }

    def build_system(self, document_sources: List[str]):
        """构建完整的 RAG 系统"""
        try:
            if isinstance(document_sources, str):
                logger.warning("document_sources 是字符串类型，转换为列表")
                document_sources = [document_sources]

            logger.info(f"开始构建RAG 系统，文档源数量: {len(document_sources)}")

            # 1. 设置嵌入模型
            self.setup_embeddings()

            # 2. 设置语言模型
            self.setup_llm()

            # 3. 加载文档
            documents = self.load_documents(document_sources)
            if not documents:
                raise ValueError("没有成功加载任何文档")

            # 4. 分割文档
            split_docs = self.split_documents(documents)

            # 5. 创建向量数据库
            self.create_vectorstore(split_docs)

            # 6. 设置问答链
            self.setup_qa_chain()

            logger.info(f"RAG 系统构建完成！使用模型: {self.model_name}")

        except Exception as e:
            logger.error(f"系统构建失败: {str(e)}")
            raise